# -*- coding: utf-8 -*-
# Author       : Xinwu
# Email        : lexinwu@genenergy.cn
# Describe     : external function for scanpy
# Created Time : 2023-06-26 16:53:34
# Last Modified: 2023-06-28 15:27:47

import scvi
import os
import logging
import scanpy as sc
import pandas as pd

logger = logging.getLogger(__name__)
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(message)s',level = logging.INFO)

def check_key(key,keys):
    """check key exists in a list or not 
    args: 
        key: which key you want to check.
        keys: a list
    """
    try:
        assert key in keys
    except AssertionError:
        logger.error(f'{key} not in {keys}, please check!')

# Doublet detection using scvi-tools
def scvi_doublet_predict(adata, batch_key, doublet_score = 1, hvg_num = 2000):
    """using scvi-tools to detection doublet
    args:
        adata: AnnData object. Rows represent cells, columns represent features.
        layer: Default: None, if not None, uses this as the key in adata.layers for raw count data.
        batch_key: key in adata.obs for batch information.
        doublet_score: doublet_score means doublet score subtract singlet score (df.doublet - df.singlet). 
                       if doublet_score greater than this value, implied that it is doublet. Default: 1
        hvg_num: Number of highly-variable genes to keep.
    return:
        return AnnData object which contain predicted doublet information in adata.obs,
        key of 'doublet' 'singlet' 'prediction' 'predicted_doublet' in adata.obs.
    """
    check_key(batch_key, adata.obs.keys())
    bdata = adata.copy()
    bdata.layers["counts"] = bdata.X.copy()
    sc.pp.normalize_total(bdata, target_sum = 1e4)
    sc.pp.log1p(bdata)
    sc.pp.highly_variable_genes(bdata, n_top_genes = hvg_num, subset = True, batch_key = batch_key)
    batch_list = list(bdata.obs[batch_key].drop_duplicates())
    bacth_key_len = len(batch_list)
    if bacth_key_len > 1:
        scvi.model.SCVI.setup_anndata(adata = bdata, batch_key = batch_key, layer = 'counts')
        vae = scvi.model.SCVI(bdata)
        vae.train()
        result_df = pd.DataFrame()
        for batch in batch_list:
            logger.info(f'Now start treate {batch}')
            solo = scvi.external.SOLO.from_scvi_model(vae, restrict_to_batch = batch)
            solo.train()
            df = solo.predict()
            df['prediction'] = solo.predict(soft = False)
            df['dif'] = df.doublet - df.singlet
            result_df = pd.concat([result_df, df], axis = 0)
        # 考虑一种特殊情况 batch_list 与 bdata.obs[batch_key] 顺序不一致, 不要觉得这种情况不会发生, 就不去考虑这种问题, 虽然可能性极低.
        # 那就需要对 result_df 按照 bdata.obs.index 进行排序, 就能够确保万无一失, 使用 result_df.reindex(bdata.obs.index).
        result_df = result_df.reindex(bdata.obs.index)
        doublets = result_df[(result_df.prediction == 'doublet') & (result_df.dif > doublet_score)]
        adata.obs['doublet'] = result_df.doublet.values
        adata.obs['singlet'] = result_df.singlet.values
        adata.obs['prediction'] = result_df.prediction.values
        adata.obs['predicted_doublet'] = adata.obs.index.isin(doublets.index)
    else:
        scvi.model.SCVI.setup_anndata(adata = bdata, layer = 'counts')
        vae = scvi.model.SCVI(bdata)
        vae.train()
        solo = scvi.external.SOLO.from_scvi_model(vae)
        solo.train()
        result_df = solo.predict()
        result_df['prediction'] = solo.predict(soft = False)
        result_df['dif'] = result_df.doublet - result_df.singlet
        doublets = result_df[(result_df.prediction == 'doublet') & (result_df.dif > doublet_score)]
        adata.obs['doublet'] = result_df.doublet.values
        adata.obs['singlet'] = result_df.singlet.values
        adata.obs['prediction'] = result_df.prediction.values
        adata.obs['predicted_doublet'] = adata.obs.index.isin(doublets.index)
    return adata

def scvi_batch_effect(adata, batch_key, layer = None,
                      categorical_covariate_keys = None,
                      continuous_covariate_keys = None):
    """using scvi-tools to remove batch effect
    args:
        adata: AnnData object. Rows represent cells, columns represent features.
        layer: Default: None, if not None, uses this as the key in adata.layers for raw count data.
        batch_key: key in adata.obs for batch information.
        categorical_covariate_keys: keys in adata.obs that correspond to categorical data(分类型变量, 以去除干扰).
        continuous_covariate_keys: keys in adata.obs that correspond to continuous data(连续型变量, 以去除干扰).
    return:
        return AnnData object which contains adata.obsm["X_scvi"] and adata.layers['scvi_normalized']
    """
    check_key(batch_key, adata.obs.keys())
    bdata = adata.copy()
    batch_list = list(bdata.obs[batch_key].drop_duplicates())
    bacth_key_len = len(batch_list)
    if bacth_key_len > 1:
        scvi.model.SCVI.setup_anndata(adata = bdata, batch_key = batch_key, layer = layer,
                                      categorical_covariate_keys = categorical_covariate_keys,
                                      continuous_covariate_keys = continuous_covariate_keys)
        vae = scvi.model.SCVI(bdata)
        vae.train()
        adata.obsm["X_scvi"] = vae.get_latent_representation()
        # shape 不一致会导致报错
        adata.layers['scvi_normalized'] = vae.get_normalized_expression(library_size = 1e4)
    else:
        scvi.model.SCVI.setup_anndata(adata = bdata, layer = layer,
                                      categorical_covariate_keys = categorical_covariate_keys,
                                      continuous_covariate_keys = continuous_covariate_keys)
        vae = scvi.model.SCVI(bdata)
        vae.train()
        adata.obsm["X_scvi"] = vae.get_latent_representation()
        # shape 不一致会导致报错
        adata.layers['scvi_normalized'] = vae.get_normalized_expression(library_size = 1e4)
    return adata, vae

def scvi_DE(adata, groupby, model = None,
            batch_key = None, layer = None,
            categorical_covariate_keys = None,
            continuous_covariate_keys = None):
    """using scvi-tools to find group marker gene
    args:
        adata: AnnData object. Rows represent cells, columns represent features.
        groupby: The key of the observations grouping to consider.
        model: where saved the scvi train model, if model is not None, will use this model to find group marker gene.
        layer: Default: None, if not None, uses this as the key in adata.layers for raw count data.
        batch_key: key in adata.obs for batch information.
        categorical_covariate_keys: keys in adata.obs that correspond to categorical data(分类型变量, 以去除干扰).
        continuous_covariate_keys: keys in adata.obs that correspond to continuous data(连续型变量, 以去除干扰).
    return:
        return AnnData object which contains adata.obsm["X_scvi"] and adata.layers['scvi_normalized']
    """
    check_key(groupby, adata.obs.keys())
    if mode is not None:
        try:
            assert os.path.exists(f'{model}/model.pt')
        except AssertionError:
            logger.error(f'{model} file not exists, please check!')
        model = scvi.model.SCVI.load(model, adata)
        full_de_res = model.differential_expression(groupby = groupby)
    else:
        check_key(batch_key, adata.obs.keys())
        bdata = adata.copy()
        batch_list = list(bdata.obs[batch_key].drop_duplicates())
        bacth_key_len = len(batch_list)
        if bacth_key_len > 1:
            scvi.model.SCVI.setup_anndata(adata = bdata, batch_key = batch_key, layer = layer,
                                          categorical_covariate_keys = categorical_covariate_keys,
                                          continuous_covariate_keys = continuous_covariate_keys)
            model = scvi.model.SCVI(bdata)
            model.train()
            full_de_res = model.differential_expression(groupby = groupby)
        else:
            scvi.model.SCVI.setup_anndata(adata = bdata, layer = layer,
                                          categorical_covariate_keys = categorical_covariate_keys,
                                          continuous_covariate_keys = continuous_covariate_keys)
            model = scvi.model.SCVI(bdata)
            model.train()
            full_de_res = model.differential_expression(groupby = groupby)
    return full_de_res